#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : zhipu.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/13 11:44
import asyncio

from apps.file_extract_util.base import BaseFileExtractUtil
from apps.utils import async_http_client


class Zhipu(BaseFileExtractUtil):

    def __init__(self, api_key, *args, **kwargs):
        self.api_key = api_key
        pass

    # 获取支持的文件类型
    async def get_support_file_type(self, *args, **kwargs):
        # .PDF .DOCX .DOC .XLS .XLSX .PPT .PPTX .PNG .JPG .JPEG .CSV .PY .TXT .MD .BMP .GIF
        return ['pdf', 'docx', 'doc', 'xls', 'xlsx', 'ppt', 'pptx', 'png', 'jpg', 'jpeg', 'csv', 'py', 'txt', 'md',
                'bmp', 'gif']

    # 文件内容提取
    async def extract_file_content(self, base_url, file_object, *args, **kwargs):
        response = await async_http_client.post(
            url=base_url,
            files={
                'file': file_object,
                'purpose': (None, 'file-extract')
            },
            headers=self._header()
        )
        print(response)
        file_id = response.get('id')
        content = await self.get_file_content(file_id=file_id)
        # 删除记录
        await self.delete_file(file_id=file_id)
        return content

    async def get_file_content(self, file_id):
        response = await async_http_client.get(
            url=f'https://open.bigmodel.cn/api/paas/v4/files/{file_id}/content',
            headers=self._header()
        )
        return response.get('content')

    async def delete_file(self, file_id):
        await async_http_client.delete(
            url=f'https://open.bigmodel.cn/api/paas/v4/files/{file_id}',
            headers=self._header()
        )

    def _header(self) -> dict:
        return {"Authorization": f"Bearer {self.api_key}"}

